--[[## ft=lua ##################################################################
  This file is licensed under [WTFPL v2](https://www.wtfpl.net/).  You may read
  the full license in the file LICENSE or by following the web link.

  FIXED - thorough fixed point operations in Lua

  REPRESENTATION

  A fixed point value is represented by a triad: `{ _v, _m, _n }`  `_v` is the
  raw  (unsigned integer) value used in calculations.  `_m` and `_n` are the
  integer and the fractional bit counts respectively.  The ARM variant of Q
  numbers is used for this.  Thus a q15.1 represents 15 bits of signed integer
  range and 1 bit of fractional range.  (In the TI variant that same value would
  be q14.1 as the sign bit is implied.)  Only Q and not UQ numbers are
  supported.

  Note that the triad is exposed only for debugging courtesy.  Manipulation of
  the values directly will result in failures both subtle and unsubtle.

  All viable operators are supported that the `number` type supports.  This
  means that all arithmetic operations are supported and all comparisons are
  supported.  These will transparently interact with the `number` type with
  said `number` being on the left or the right of the operator with one
  exception: `==`.  (The semantics of equality comparison in Lua and how they
  relate to metamethods mean that only two fixed numbers can be compared for
  equality.)

  A lot of extra operations are also provided.  Consult the README.md file for
  full details.

  There are two levels of constructor in this library.  The first one, returned
  by `require 'fixed'` will generate value factories.  It stores the desired
  `m` and `n` factors for numbers and each time it is called it will return an
  actual fixed number of that type.  It is typically used something like this:

    fixed = require 'fixed'
    -- a q15.1 number
    q15_1 = fixed(15, 1)

  The second level of constructor is a value constructor.  It is what the value
  factories use to create actual fixed values:

    fixed = require 'fixed'
    q15_1 = fixed(15, 1)
    a = q15_1(1.5)
    b = q15_1(2.2)

  All fixed numbers are interoperable, no matter what the `m` and `n` values
  given.  Conversions between them, however, will lose precision.

--############################################################################]]

local math, string = math, string

--[[###################################
  Caches and numbers.
--###################################]]
local qcache = {}
local vmcache = {}
local tcache = {}
local MAX_BITS = 63

--[[###################################
  Module-local helper functions
--###################################]]

-- generate an n-bit mask
local function mask(bits)
    return (1 << bits) - 1
end

-- convert an unsigned value of given bit width to signed
local function to_signed(value, bits)
    if value >= (1 << (bits - 1)) then
        return value - (1 << bits)
    else
        return value
    end
end

-- convert a signed value of given bit width to unsigned
local function from_signed(value, bits)
    -- Convert signed value to unsigned value
    if value < 0 then
        return value + (1 << bits)
    else
        return value
    end
end

-- convert fixed value to double (Lua-native) value
local function fixed_to_double(q)
    local signed_value = to_signed(q._v, q._m + q._n)
    return (signed_value * 1.0) / (1 << q._n)
end

-- convert double (Lua-native) value to fixed value
local function double_to_fixed(d, n)
    return math.floor(d * (1 << n) + 0.5)
end

local function check_fixed_type(q)
    return type(q) == 'table' and
           q._v ~= nil and math.type(q._v) == 'integer' and
           q._m ~= nil and math.type(q._m) == 'integer' and
           q._n ~= nil and math.type(q._n) == 'integer' and
           string.match(getmetatable(q)['__type__'],'^q%d+%.%d+$')
end

local function normalize_numbers(v1, v2, err_msg)
    if type(v1) == 'number' then
        local v = double_to_fixed(v1, v2._n)
        v1 = v2:clone()
        v1._v = v
    elseif type(v2) == 'number' then
        local v = double_to_fixed(v2, v1._n)
        v2 = v1:clone()
        v2._v = v
    elseif not check_fixed_type(v2) then
        error(err_msg)
    end
    return v1, v2
end

local function match_numbers(q1, q2)
    local r = q1:clone()
    if q1._n ~= q2._n then
        local common_n = math.max(q1._n, q2._n)
        local temp = q2:clone()
        temp:rescale(q2._m, common_n)
        q2 = temp
        r:rescale(q1._m, common_n)
    end
    return r, q2
end

local function build_value_metatable(m, n, type_str)
    return {
        -- OPERATIONS
        __index = {
            -- type checking operations
            get_type = function(q)
                return type_str
            end,

            -- output and display operations
            printv = function(q)
                print(string.format('m = %d, n = %d, v = %x -> %f', q._m, q._n, q._v, q()))
            end,

            serialize = function(q)
                return string.format('fixed:%d:%d:%d', q._m, q._n, q._v)
            end,

            to_hex = function(q)
                return string.format('0x%x', q._v)
            end,

            to_string = function(q)
                return string.format('q%d.%d(%f)', q._m, q._n, fixed_to_double(q))
            end,

            -- format manipulation operations
            clone = function(q)
                return setmetatable({
                    _v = q._v,
                    _m = q._m,
                    _n = q._n }, getmetatable(q))
            end,

            mask = function(q)
                q._v = q._v & mask(q._m + q._n)
            end,

            rescale = function(q, new_m, new_n)
                local v = q._v
                local m = q._m
                local n = q._n
                local total_bits = m + n

                -- Adjust fractional bits
                if new_n > n then
                    v = v << (new_n - n)
                elseif new_n < n then
                    -- Round when reducing fractional bits
                    local round = 1 << (n - new_n - 1)
                    v = (v + round) >> (n - new_n)
                end

                -- Adjust integer bits
                local new_total = new_m + new_n
                if new_total > total_bits then
                    -- Extend sign if needed
                    local sign_bit = v & (1 << (total_bits - 1))
                    if sign_bit ~= 0 then
                        local extension = mask(new_total - total_bits)
                        v = v | (extension << total_bits)
                    end
                elseif new_total < total_bits then
                    -- Mask to new size
                    v = v & mask(new_total)
                end

                q._v = v
                q._m = new_m
                q._n = new_n
            end,

            saturate = function(q)
                local total_bits = q._m + q._n
                local max_unsigned = mask(total_bits)
                local raw_max = mask(total_bits - 1)

                if q._v > max_unsigned then
                    q._v = raw_max
                end

                return q
            end,

            -- predicate tests
            check_overflow = function(q)
                local total_bits = q._m + q._n
                local value = q._v

                -- Check if value exceeds the representable range
                if (value & ~mask(total_bits)) ~= 0 then
                    return true
                end

                return false
            end,

            is_zero = function(q)
                return q._v == 0
            end,

            is_positive = function(q)
                local signed_value = to_signed(q._v, q._m + q._n)
                return signed_value > 0
            end,

            is_negative = function(q)
                local signed_value = to_signed(q._v, q._m + q._n)
                return signed_value < 0
            end,

            -- arithmetic operations (non-operator)
            abs = function(q)
                local r = q:clone()
                local signed_value = to_signed(r._v, r._m + r._n)
                r._v = from_signed(math.abs(signed_value), r._m + r._n)
                return r
            end,

            -- conversion operations
            to_number = function(q)
                return fixed_to_double(q)
            end,

            to_integer = function(q)
                local signed_value = to_signed(q._v, q._m + q._n)
                return math.floor(signed_value / (1 << q._n))
            end,

            to_fraction = function(q)
                return q._v & mask(q._n)
            end,
        },

        -- OPERATORS
        __add = function(q1, q2)
            q1, q2 = normalize_numbers(q1, q2, 'attempt to add fixed-point with non-numeric value')
            q1, q2 = match_numbers(q1, q2)

            local signed_v1 = to_signed(q1._v, q1._m + q1._n)
            local signed_v2 = to_signed(q2._v, q2._m + q2._n)
            local result = signed_v1 + signed_v2
            q1._v = from_signed(result, q1._m + q1._n)

            q1:mask()
            if q1:check_overflow() then
                q1:saturate()
            end
            return q1
        end,

        __sub = function(q1, q2)
            q1, q2 = normalize_numbers(q1, q2, 'attempt to subtract fixed-point with non-numeric value')
            q1, q2 = match_numbers(q1, q2)

            local signed_v1 = to_signed(q1._v, q1._m + q1._n)
            local signed_v2 = to_signed(q2._v, q2._m + q2._n)
            local result = signed_v1 - signed_v2
            q1._v = from_signed(result, q1._m + q1._n)

            q1:mask()
            if q1:check_overflow() then
                q1:saturate()
            end
            return q1
        end,

        __mul = function(q1, q2)
            q1, q2 = normalize_numbers(q1, q2, 'attempt to multiply fixed-point with non-numeric value')
            q1, q2 = match_numbers(q1, q2)

            -- Convert to signed values for arithmetic
            local signed_v1 = to_signed(q1._v, q1._m + q1._n)
            local signed_v2 = to_signed(q2._v, q2._m + q2._n)

            -- Multiply and then shift right by n to maintain scale
            local intermediate = signed_v1 * signed_v2
            local result = intermediate >> q1._n
            q1._v = from_signed(result, q1._m + q1._n)

            q1:mask()
            if q1:check_overflow() then
                q1:saturate()
            end
            return q1
        end,

        __div = function(q1, q2)
            q1, q2 = normalize_numbers(q1, q2, 'attempt to divide fixed-point with non-numeric value')
            q1, q2 = match_numbers(q1, q2)

            local result = fixed_to_double(q1) / fixed_to_double(q2)
            q1._v = double_to_fixed(result, q1._n)

            q1:mask()
            if q1:check_overflow() then
                q1:saturate()
            end
            return q1
        end,

        __mod = function(q1, q2)
            q1, q2 = normalize_numbers(q1, q2, 'attempt to mod fixed-point with non-numeric value')
            q1, q2 = match_numbers(q1, q2)

            local signed_v1 = to_signed(q1._v, q1._m + q1._n)
            local signed_v2 = to_signed(q2._v, q2._m + q2._n)
            local result = signed_v1 % signed_v2
            q1._v = from_signed(result, q1._m + q1._n)

            q1:mask()
            return q1
        end,

        __pow = function(q1, q2)
            if (not check_fixed_type(q1) and type(q1) ~= 'number') or
               (not check_fixed_type(q2) and type(q2) ~= 'number') then
                error 'attempt to pow fixed-point with non-numeric value'
            end
            local r, v1, v2
            if check_fixed_type(q1) and check_fixed_type(q2) then
                q1, q2 = match_numbers(q1, q2)
            end
            if check_fixed_type(q1) then
                r = r or q1:clone()
                v1 = fixed_to_double(q1)
            else
                v1 = q1
            end
            if check_fixed_type(q2) then
                r = r or q2:clone()
                v2 = fixed_to_double(q2)
            else
                v2 = q2
            end

            -- Use floating point for non-integer exponents
            local result = v1 ^ v2
            r._v = double_to_fixed(result, r._n)
            r:mask()
            return r
        end,

        __unm = function(q)
            local r = q:clone()
            local signed_value = to_signed(r._v, r._m + r._n)
            r._v = from_signed(-signed_value, r._m + r._n)
            r:mask()
            if r:check_overflow() then
                r:saturate()
            end
            return r
        end,

        __concat = function(a, b)
            return tostring(a) .. tostring(b)
        end,

        __eq = function(q1, q2)
            if not check_fixed_type(q2) then
                return false
            end
            q1, q2 = match_numbers(q1, q2)
            return q1._v == q2._v
        end,

        __lt = function(q1, q2)
            q1, q2 = normalize_numbers(q1, q2, 'attempt to compare fixed-point with non-numeric value')
            q1, q2 = match_numbers(q1, q2)

            local signed_v1 = to_signed(q1._v, q1._m + q1._n)
            local signed_v2 = to_signed(q2._v, q2._m + q2._n)
            return signed_v1 < signed_v2
        end,

        __le = function(q1, q2)
            q1, q2 = normalize_numbers(q1, q2, 'attempt to compare fixed-point with non-numeric value')
            q1, q2 = match_numbers(q1, q2)

            local signed_v1 = to_signed(q1._v, q1._m + q1._n)
            local signed_v2 = to_signed(q2._v, q2._m + q2._n)
            return signed_v1 <= signed_v2
        end,

        __call = function(q)
            return fixed_to_double(q)
        end,

        __tostring = function(q)
            return string.format('q%d.%d(%f)', q._m, q._n, fixed_to_double(q))
        end,

        -- TYPE IDENTITY
        __type__ = type_str,
    }
end

local function build_constructor_methods(type_str)
    return {
        zero = function(self)
            if not tcache[type_str].zero then tcache[type_str].zero = self(0) end
            return tcache[type_str].zero
        end,

        one = function(self)
            if not tcache[type_str].one then tcache[type_str].one = self(1) end
            return tcache[type_str].one
        end,

        max_value = function(self)
            if not tcache[type_str].max_value then
                local raw = mask(self._m + self._n - 1)
                tcache[type_str].max_value = self:from_raw(raw)
            end
            return tcache[type_str].max_value
        end,

        min_value = function(self)
            if not tcache[type_str].min_value then
                local raw = 1 << (self._m + self._n - 1)
                tcache[type_str].min_value = self:from_raw(raw)
            end
            return tcache[type_str].min_value
        end,

        from_raw = function(self, raw_value)
            local instance = self(0)
            instance._v = raw_value
            if instance:check_overflow() then
                error(string.format('raw_value = %d, max_value = %d', raw_value, self:max_value()._v))
            else
                return instance
            end
        end,

        get_precision = function(self)
            return self._m, self._n
        end,

        from_serialized = function(self, str)
            local m, n, v = str:match('fixed:(%d+):(%d+):(%d+)')
            if not m then error('Invalid serialized format') end
            m, n, v = tonumber(m), tonumber(n), tonumber(v)
            if m ~= self._m or n ~= self._n then error('Type mismatch in deserialization') end
            return setmetatable({_v = v, _m = m, _n = n}, getmetatable(self(0)))
        end,
    }
end

--[[###################################
  Builder constructor.
  'q15.1' in ARM notation is called `new(15, 1)`
  ###################################]]
local function new(_, m, n)

    assert(math.type(m) == 'integer' and math.type(n) == 'integer', 'm and n must be integers')
    assert(m >= 1, 'm (integer bits) must be at least 1 (for sign bit)')
    assert(n >= 0, 'n (fractional bits) must be non-negative')
    assert(m + n < MAX_BITS, string.format('Total bits must be less than %d (got %d+%d=%d)', MAX_BITS, m, n, m+n))

    local _m = m
    local _n = n
    local total_bits = m + n
    local type_str = string.format('q%d.%d', _m, _n)

    -- value constructor
    local function newv(q, v)
        local raw_value
        if type(v) == 'number' then
            raw_value = double_to_fixed(v, _n)
        elseif check_fixed_type(v) then
            -- Direct conversion between different fixed-point formats
            if _n == v._n then
                raw_value = v._v
            else
                local shift = _n - v._n
                if shift > 0 then
                    raw_value = v._v << shift
                else
                    -- Round when reducing fractional bits
                    local round = 1 << (-shift - 1)
                    raw_value = (v._v + round) >> (-shift)
                end
            end
        else
            error('Invalid value type: ' .. type(v))
        end

        -- Create the metatable for the value based on the type string
        if not vmcache[type_str] then
            vmcache[type_str] = build_value_metatable(_m, _n, type_str)
        end

        local r = setmetatable({
            _v = raw_value,
            _m = _m,
            _n = _n,
        }, vmcache[type_str])

        r:mask()
        return r
    end

    if not qcache[type_str] then
        tcache[type_str] = {}

        -- Create the type constructor with additional methods
        local constructor_methods = build_constructor_methods(type_str)
        local constructor = setmetatable({
            _m = _m,
            _n = _n,
        }, {
            __call = newv,
            __type__ = type_str,
            __index = constructor_methods,
        })
        qcache[type_str] = constructor
    end

    return qcache[type_str]
end

return setmetatable({}, { __call = new })
